/* Copyright (c) 2008, John Cronin
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the copyright holder nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */

using System;
using System.Collections.Generic;
using System.Text;
using System.IO;

namespace ElfLib
{
    public class Elf64Writer : tysila.IOutputFile
    {
        Elf64_Ehdr ehdr;
        string e_point;

        public void SetEntryPoint(string entry_point)
        { e_point = entry_point; }

        public Elf64Writer(string source_filename)
        {
            ehdr = new Elf64_Ehdr(source_filename);
        }

        private bool WriteObjectFile(System.IO.Stream output)
        {
            // Write out the elf file.

            // Relocatable files are:
            //
            //   ELF header
            //   Program header table (optional)
            //   Section 1
            //   Section 2
            //   ...
            //   Section n
            //   Section header table

            // Find the entry point symbol
            if (e_point != null)
            {
                Elf64_Symbol_Shdr.Elf64_Sym entry_sym = null;
                Elf64_Symbol_Shdr.Elf64_Sym epoint_sym = null;
                foreach (Elf64_Symbol_Shdr.Elf64_Sym sym in ehdr.e_syms.defined_syms)
                {
                    if (sym.name == "_start")
                        entry_sym = sym;
                    if (sym.name == e_point)
                        epoint_sym = sym;
                }

                if (entry_sym != null)
                    ehdr.e_syms.defined_syms.Remove(entry_sym);
                if (epoint_sym != null)
                    AddTextSymbol((int)epoint_sym.st_value, "_start");
            }

            // First, determine the offsets of the various tables
            ulong cur_offset;

            if (ehdr.phdrs.Count > 0)
            {
                ehdr.e_phoff = (ulong)Align(Elf64_Ehdr.GetLength(), 8);
                ehdr.e_phentsize = (ushort)Elf64_Phdr.GetLength();
                cur_offset = ehdr.e_phoff + (ulong)ehdr.e_phnum * (ulong)ehdr.e_phentsize;}
            else
            {
                ehdr.e_phoff = 0;
                ehdr.e_phentsize = 0;
                cur_offset = (ulong)Align(Elf64_Ehdr.GetLength(), 8);
            }            

            ehdr.GenerateRelocationSymbols();

            foreach (Elf64_Shdr shdr in ehdr.shdrs)
            {
                if (shdr.sh_type != SectionType.SHT_NULL)
                {
                    shdr.sh_offset = Align(cur_offset, 8);
                    cur_offset = shdr.sh_offset + shdr.sh_size;
                }
            }

            ehdr.e_shoff = (ulong)Align(cur_offset, 8);

            // Now write the actual elf header table
            output.Seek(0, SeekOrigin.Begin);
            ehdr.Write(output);

            // Write program header table
            output.Seek((long)ehdr.e_phoff, SeekOrigin.Begin);
            foreach (Elf64_Phdr phdr in ehdr.phdrs)
                phdr.Write(output);

            // Write actual section data
            foreach (Elf64_Shdr shdr in ehdr.shdrs)
            {
                output.Seek((long)shdr.sh_offset, SeekOrigin.Begin);
                Write(output, shdr.data);
            }

            // Write section header data
            output.Seek((long)ehdr.e_shoff, SeekOrigin.Begin);
            foreach (Elf64_Shdr shdr in ehdr.shdrs)
                shdr.Write(output);

            return true;
        }

        public void AlignText(int a)
        { Align(ehdr.text.data, a); }
        public void AlignData(int a)
        { Align(ehdr.data.data, a); }
        public void AlignRodata(int a)
        { Align(ehdr.rodata.data, a); }

        private void Align(IList<byte> iList, int a)
        {
            while ((iList.Count % a) != 0)
                iList.Add(0);
        }

        private int Align(int v, int align)
        {
            int extra = v % align;
            if (extra == 0)
                return v;
            else
                return v - extra + align;
        }

        private ulong Align(ulong v, ulong align)
        {
            ulong extra = v % align;
            if (extra == 0)
                return v;
            else
                return v - extra + align;
        }

        public static void Write(object s, byte v)
        {
            if (s is Stream)
                ((Stream)s).WriteByte(v);
            else if (s is ICollection<byte>)
                ((ICollection<byte>)s).Add(v);
            else
                throw new Exception("Unknown container type");
        }

        public static void Write(object s, UInt16 v)
        {
            Write(s, (byte)(v & 0xff));
            Write(s, (byte)((v >> 8) & 0xff));
        }

        public static void Write(object s, UInt32 v)
        {
            Write(s, (UInt16)(v & 0xffff));
            Write(s, (UInt16)((v >> 16) & 0xffff));
        }

        public static void Write(object s, UInt64 v)
        {
            Write(s, (UInt32)(v & 0xffffffff));
            Write(s, (UInt32)((v >> 32) & 0xffffffff));
        }

        public static void Write(object s, Int64 v)
        {
            if (v >= 0)
                Write(s, (UInt64)v);
            else
                Write(s, (~((UInt64)(-v))) + 1);
        }

        public static void Write(object s, IEnumerable<byte> v)
        {
            if (v == null)
                return;
            foreach (byte b in v)
                Write(s, b);
        }

        #region IOutputFile Members

        public IList<byte> GetText()
        {
            return ehdr.text.data;
        }

        public IList<byte> GetData()
        {
            return ehdr.data.data;
        }

        public IList<byte> GetRodata()
        {
            return ehdr.rodata.data;
        }

        public void AddTextSymbol(int offset, string name)
        {
            ehdr.e_syms.defined_syms.Add(new Elf64_Symbol_Shdr.Elf64_Sym
            {
                name = name,
                st_name = ehdr.e_symstr.GetOffset(name),
                st_info = Elf64_Symbol_Shdr.Elf64_Sym.BindingFlags.STB_GLOBAL |
                Elf64_Symbol_Shdr.Elf64_Sym.SymbolTypes.STT_FUNC,
                st_shndx = ehdr.text.index,
                st_value = Convert.ToUInt64(offset),
                st_size = 0,
                st_other = 0
            });
        }

        public void AddDataSymbol(int offset, string name)
        {
            ehdr.e_syms.defined_syms.Add(new Elf64_Symbol_Shdr.Elf64_Sym
            {
                name = name,
                st_name = ehdr.e_symstr.GetOffset(name),
                st_info = Elf64_Symbol_Shdr.Elf64_Sym.BindingFlags.STB_GLOBAL |
                Elf64_Symbol_Shdr.Elf64_Sym.SymbolTypes.STT_FUNC,
                st_shndx = ehdr.data.index,
                st_value = Convert.ToUInt64(offset),
                st_size = 0,
                st_other = 0
            });
        }

        public void AddRodataSymbol(int offset, string name)
        {
            ehdr.e_syms.defined_syms.Add(new Elf64_Symbol_Shdr.Elf64_Sym
            {
                name = name,
                st_name = ehdr.e_symstr.GetOffset(name),
                st_info = Elf64_Symbol_Shdr.Elf64_Sym.BindingFlags.STB_GLOBAL |
                Elf64_Symbol_Shdr.Elf64_Sym.SymbolTypes.STT_FUNC,
                st_shndx = ehdr.rodata.index,
                st_value = Convert.ToUInt64(offset),
                st_size = 0,
                st_other = 0
            });
        }

        public void AddTextRelocation(int offset, string name, uint rel_type, int value)
        {
            ehdr.relatext.relocs.Add(new Elf64_Rela_Shdr.Elf64_Rela((ulong)offset, rel_type, value, name));
        }

        public void AddDataRelocation(int offset, string name, uint rel_type, int value)
        {
            ehdr.reladata.relocs.Add(new Elf64_Rela_Shdr.Elf64_Rela((ulong)offset, rel_type, value, name));
        }

        public void AddRodataRelocation(int offset, string name, uint rel_type, int value)
        {
            ehdr.relarodata.relocs.Add(new Elf64_Rela_Shdr.Elf64_Rela((ulong)offset, rel_type, value, name));
        }

        public void Write(Stream output)
        {
            WriteObjectFile(output);
        }

        #endregion
    }
}
